#!/usr/bin/env python3
"""
EndoSight-UC V4.0 Security Audit Script
Automated security vulnerability scanning and dependency checking
"""

import subprocess
import sys
import json
import re
from pathlib import Path
from typing import List, Dict, Any
import structlog

logger = structlog.get_logger()


class SecurityAuditor:
    """Security vulnerability auditor for dependencies and code"""

    def __init__(self):
        self.vulnerabilities = []
        self.warnings = []

    def check_dependencies_vulnerabilities(self) -> Dict[str, Any]:
        """Check for known vulnerabilities in dependencies"""
        logger.info("开始检查依赖库安全漏洞...")

        try:
            # Try to use pip-audit if available
            result = subprocess.run(
                [sys.executable, "-m", "pip_audit", "--format", "json"],
                capture_output=True,
                text=True,
                cwd=Path.cwd()
            )

            if result.returncode == 0:
                try:
                    audit_data = json.loads(result.stdout)
                    vulnerabilities = audit_data.get("vulnerabilities", [])

                    logger.info(f"发现 {len(vulnerabilities)} 个已知漏洞")
                    return {
                        "status": "completed",
                        "vulnerabilities": vulnerabilities,
                        "total_count": len(vulnerabilities)
                    }
                except json.JSONDecodeError:
                    logger.warning("pip-audit 输出格式解析失败")
            else:
                logger.warning("pip-audit 不可用，使用基础检查")

        except Exception as e:
            logger.error(f"依赖漏洞检查失败: {e}")

        # Fallback: Basic version checking for critical packages
        return self._basic_dependency_check()

    def _basic_dependency_check(self) -> Dict[str, Any]:
        """Basic dependency version checking"""
        critical_packages = {
            'fastapi': '0.104.1',  # Below this may have security issues
            'uvicorn': '0.24.0',
            'pydantic': '2.5.0',
            'torch': '2.1.0',
            'transformers': '4.36.0'
        }

        vulnerable_packages = []

        try:
            import pkg_resources
            for package, min_version in critical_packages.items():
                try:
                    installed_version = pkg_resources.get_distribution(package).version
                    if self._compare_versions(installed_version, min_version) < 0:
                        vulnerable_packages.append({
                            'name': package,
                            'installed_version': installed_version,
                            'required_version': min_version,
                            'severity': 'HIGH'
                        })
                except pkg_resources.DistributionNotFound:
                    vulnerable_packages.append({
                        'name': package,
                        'installed_version': 'NOT_INSTALLED',
                        'required_version': min_version,
                        'severity': 'CRITICAL'
                    })

        except Exception as e:
            logger.error(f"基础依赖检查失败: {e}")

        return {
            "status": "basic_check_completed",
            "vulnerabilities": vulnerable_packages,
            "total_count": len(vulnerable_packages)
        }

    def _compare_versions(self, v1: str, v2: str) -> int:
        """Compare version strings (simple implementation)"""
        def normalize(v):
            return [int(x) for x in re.sub(r'[^0-9.]', '', v).split('.')]

        v1_parts = normalize(v1)
        v2_parts = normalize(v2)

        # Pad shorter version with zeros
        max_len = max(len(v1_parts), len(v2_parts))
        v1_parts.extend([0] * (max_len - len(v1_parts)))
        v2_parts.extend([0] * (max_len - len(v2_parts)))

        for a, b in zip(v1_parts, v2_parts):
            if a < b:
                return -1
            elif a > b:
                return 1
        return 0

    def scan_code_for_secrets(self) -> List[Dict[str, Any]]:
        """Scan code for hardcoded secrets and sensitive information"""
        logger.info("扫描代码中的硬编码敏感信息...")

        secret_patterns = {
            'api_key': [
                r'api_key\s*=\s*["\'][^"\']+["\']',
                r'API_KEY\s*=\s*["\'][^"\']+["\']',
                r'secret_key\s*=\s*["\'][^"\']+["\']',
                r'SECRET_KEY\s*=\s*["\'][^"\']+["\']'
            ],
            'password': [
                r'password\s*=\s*["\'][^"\']+["\']',
                r'PASSWORD\s*=\s*["\'][^"\']+["\']'
            ],
            'connection_string': [
                r'mongodb://[^@]+:[^@]+@',
                r'mysql://[^@]+:[^@]+@',
                r'postgres://[^@]+:[^@]+@'
            ],
            'private_key': [
                r'-----BEGIN [A-Z]+ KEY-----',
                r'-----BEGIN RSA PRIVATE KEY-----'
            ]
        }

        secrets_found = []

        # Scan Python files
        for py_file in Path.cwd().rglob("*.py"):
            if 'venv' in str(py_file) or '__pycache__' in str(py_file):
                continue

            try:
                with open(py_file, 'r', encoding='utf-8', errors='ignore') as f:
                    content = f.read()

                for secret_type, patterns in secret_patterns.items():
                    for pattern in patterns:
                        matches = re.finditer(pattern, content, re.IGNORECASE)
                        for match in matches:
                            line_number = content[:match.start()].count('\n') + 1
                            secrets_found.append({
                                'type': secret_type,
                                'file': str(py_file),
                                'line': line_number,
                                'match': match.group()[:50] + '...' if len(match.group()) > 50 else match.group(),
                                'severity': 'HIGH' if secret_type in ['api_key', 'private_key'] else 'MEDIUM'
                            })

            except Exception as e:
                logger.error(f"扫描文件失败 {py_file}: {e}")

        logger.info(f"发现 {len(secrets_found)} 个潜在敏感信息")
        return secrets_found

    def check_file_permissions(self) -> List[Dict[str, Any]]:
        """Check file permissions for security issues"""
        logger.info("检查文件权限...")

        permission_issues = []

        # Check sensitive files
        sensitive_files = [
            '.env',
            '.env.example',
            'core/config.py',
            'requirements.txt',
            'Dockerfile'
        ]

        for file_path in sensitive_files:
            full_path = Path.cwd() / file_path
            if full_path.exists():
                stat_info = full_path.stat()
                mode = oct(stat_info.st_mode)[-3:]

                # Check if file is world-writable
                if mode[-1] in ['2', '3', '6', '7']:  # World writable
                    permission_issues.append({
                        'file': file_path,
                        'permissions': mode,
                        'issue': 'World-writable permissions',
                        'severity': 'HIGH'
                    })

                # Check if file is world-readable for sensitive files
                if file_path.startswith('.env') and mode[-1] in ['4', '5', '6', '7']:  # World readable
                    permission_issues.append({
                        'file': file_path,
                        'permissions': mode,
                        'issue': 'World-readable sensitive file',
                        'severity': 'MEDIUM'
                    })

        logger.info(f"发现 {len(permission_issues)} 个权限问题")
        return permission_issues

    def check_docker_security(self) -> List[Dict[str, Any]]:
        """Check Dockerfile for security issues"""
        logger.info("检查Docker安全配置...")

        dockerfile_path = Path.cwd() / "Dockerfile"
        security_issues = []

        if not dockerfile_path.exists():
            return [{'issue': 'Dockerfile not found', 'severity': 'MEDIUM'}]

        try:
            with open(dockerfile_path, 'r') as f:
                dockerfile_content = f.readlines()

            # Check for running as root
            root_user_patterns = ['USER root', 'USER 0']
            for i, line in enumerate(dockerfile_content, 1):
                line = line.strip()
                if any(pattern in line for pattern in root_user_patterns):
                    security_issues.append({
                        'line': i,
                        'content': line,
                        'issue': 'Running as root user',
                        'severity': 'HIGH'
                    })

                # Check for adding all files
                if line.startswith('ADD .') or line.startswith('COPY .'):
                    security_issues.append({
                        'line': i,
                        'content': line,
                        'issue': 'Adding all files to container',
                        'severity': 'MEDIUM'
                    })

        except Exception as e:
            logger.error(f"Dockerfile检查失败: {e}")

        logger.info(f"发现 {len(security_issues)} 个Docker安全问题")
        return security_issues

    def generate_security_report(self) -> Dict[str, Any]:
        """Generate comprehensive security report"""
        logger.info("生成安全审计报告...")

        report = {
            'timestamp': str(Path.cwd()),
            'scan_results': {
                'dependencies': self.check_dependencies_vulnerabilities(),
                'secrets': self.scan_code_for_secrets(),
                'file_permissions': self.check_file_permissions(),
                'docker_security': self.check_docker_security()
            },
            'summary': {}
        }

        # Calculate summary
        total_issues = 0
        critical_count = 0
        high_count = 0

        for category, results in report['scan_results'].items():
            if isinstance(results, dict) and 'vulnerabilities' in results:
                issues = results['vulnerabilities']
            else:
                issues = results

            total_issues += len(issues)

            for issue in issues:
                severity = issue.get('severity', 'UNKNOWN').upper()
                if severity == 'CRITICAL':
                    critical_count += 1
                elif severity == 'HIGH':
                    high_count += 1

        report['summary'] = {
            'total_issues': total_issues,
            'critical_issues': critical_count,
            'high_issues': high_count,
            'security_score': max(0, 10 - critical_count - high_count // 2)
        }

        return report

    def save_report(self, report: Dict[str, Any], filename: str = "security_audit_report.json"):
        """Save security report to file"""
        report_path = Path.cwd() / filename
        try:
            with open(report_path, 'w', encoding='utf-8') as f:
                json.dump(report, f, indent=2, ensure_ascii=False)
            logger.info(f"安全报告已保存: {report_path}")
        except Exception as e:
            logger.error(f"保存报告失败: {e}")


def main():
    """Main security audit function"""
    print("🔒 EndoSight-UC V4.0 安全审计工具")
    print("=" * 50)

    auditor = SecurityAuditor()
    report = auditor.generate_security_report()
    auditor.save_report(report)

    # Print summary
    summary = report['summary']
    print(f"\n📊 安全审计总结:")
    print(f"   总问题数: {summary['total_issues']}")
    print(f"   严重问题: {summary['critical_issues']}")
    print(f"   高危问题: {summary['high_issues']}")
    print(f"   安全评分: {summary['security_score']}/10")

    if summary['critical_issues'] > 0:
        print(f"\n🚨 发现严重安全问题，请立即修复！")
        return 1
    elif summary['high_issues'] > 0:
        print(f"\n⚠️  发现高危问题，建议尽快修复")
        return 1
    else:
        print(f"\n✅ 安全检查通过")
        return 0


if __name__ == "__main__":
    exit_code = main()
    sys.exit(exit_code)